package com.winter.common.dynamic.datasource;

import com.baomidou.dynamic.datasource.AbstractRoutingDataSource;
import com.baomidou.dynamic.datasource.DynamicGroupDataSource;
import com.baomidou.dynamic.datasource.provider.DynamicDataSourceProvider;
import com.baomidou.dynamic.datasource.strategy.DynamicDataSourceStrategy;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.p6spy.engine.spy.P6DataSource;
import io.seata.rm.datasource.DataSourceProxy;
import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@Data
public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean, DisposableBean {
    private static final Logger log = LoggerFactory.getLogger(DynamicRoutingDataSource.class);
    private static final String UNDERLINE = "_";
    private DynamicDataSourceProvider provider;
    private Class<? extends DynamicDataSourceStrategy> strategy;
    private String primary;
    private boolean strict;
    private boolean p6spy;
    private boolean seata;
    private Map<String, DataSource> dataSourceMap = new LinkedHashMap();
    private Map<String, DynamicGroupDataSource> groupDataSources = new ConcurrentHashMap();

    public DynamicRoutingDataSource() {
    }

    public DataSource determineDataSource() {
        return this.getDataSource(DynamicDataSourceContextHolder.peek());
    }

    private DataSource determinePrimaryDataSource() {
        log.debug("dynamic-datasource switch to the primary datasource");
        return this.groupDataSources.containsKey(this.primary) ? ((DynamicGroupDataSource)this.groupDataSources.get(this.primary)).determineDataSource() : (DataSource)this.dataSourceMap.get(this.primary);
    }

    public Map<String, DataSource> getCurrentDataSources() {
        return this.dataSourceMap;
    }

    public Map<String, DynamicGroupDataSource> getCurrentGroupDataSources() {
        return this.groupDataSources;
    }

    public DataSource getDataSource(String ds) {
        if (StringUtils.isEmpty(ds)) {
            return this.determinePrimaryDataSource();
        } else if (!this.groupDataSources.isEmpty() && this.groupDataSources.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            return ((DynamicGroupDataSource)this.groupDataSources.get(ds)).determineDataSource();
        } else if (this.dataSourceMap.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            return (DataSource)this.dataSourceMap.get(ds);
        } else if (this.strict) {
            throw new RuntimeException("dynamic-datasource could not find a datasource named" + ds);
        } else {
            return this.determinePrimaryDataSource();
        }
    }

    public synchronized void addDataSource(String ds, DataSource dataSource) {
        if (this.p6spy) {
            dataSource = new P6DataSource((DataSource)dataSource);
            log.info("dynamic-datasource [{}] wrap p6spy plugin", ds);
        }

        if (this.seata) {
            dataSource = new DataSourceProxy((DataSource)dataSource);
            log.info("dynamic-datasource [{}] wrap seata plugin", ds);
        }

        this.dataSourceMap.put(ds, dataSource);
        if (ds.contains("_")) {
            String group = ds.split("_")[0];
            if (this.groupDataSources.containsKey(group)) {
                ((DynamicGroupDataSource)this.groupDataSources.get(group)).addDatasource((DataSource)dataSource);
            } else {
                try {
                    DynamicGroupDataSource groupDatasource = new DynamicGroupDataSource(group, (DynamicDataSourceStrategy)this.strategy.newInstance());
                    groupDatasource.addDatasource((DataSource)dataSource);
                    this.groupDataSources.put(group, groupDatasource);
                } catch (Exception var5) {
                    log.error("dynamic-datasource - add the datasource named [{}] error", ds, var5);
                    this.dataSourceMap.remove(ds);
                }
            }
        }

        log.info("dynamic-datasource - load a datasource named [{}] success", ds);
    }

    public synchronized void removeDataSource(String ds) {
        if (this.dataSourceMap.containsKey(ds)) {
            DataSource dataSource = (DataSource)this.dataSourceMap.get(ds);
            this.dataSourceMap.remove(ds);
            if (ds.contains("_")) {
                String group = ds.split("_")[0];
                if (this.groupDataSources.containsKey(group)) {
                    ((DynamicGroupDataSource)this.groupDataSources.get(group)).removeDatasource(dataSource);
                }
            }

            log.info("dynamic-datasource - remove the database named [{}] success", ds);
        } else {
            log.warn("dynamic-datasource - could not find a database named [{}]", ds);
        }

    }

    public void setP6spy(boolean p6spy) {
        if (p6spy) {
            try {
                Class.forName("com.p6spy.engine.spy.P6DataSource");
                log.info("dynamic-datasource detect P6SPY plugin and enabled it");
                this.p6spy = true;
            } catch (Exception var3) {
                log.warn("dynamic-datasource enabled P6SPY ,however without p6spy dependency");
            }
        } else {
            this.p6spy = false;
        }

    }

    public void setSeata(boolean seata) {
        if (seata) {
            try {
                Class.forName("io.seata.rm.datasource.DataSourceProxy");
                this.seata = true;
                log.info("dynamic-datasource detect ALIBABA SEATA and enabled it");
            } catch (Exception var3) {
                this.seata = false;
                log.warn("dynamic-datasource enabled ALIBABA SEATA  ,however without seata dependency");
            }
        }

    }

    public void destroy() throws Exception {
        log.info("dynamic-datasource start closing ....");
        Iterator var1 = this.dataSourceMap.entrySet().iterator();

        while(var1.hasNext()) {
            Map.Entry<String, DataSource> item = (Map.Entry)var1.next();
            DataSource dataSource = (DataSource)item.getValue();
            if (this.seata) {
                DataSourceProxy dataSourceProxy = (DataSourceProxy)dataSource;
                dataSource = dataSourceProxy.getTargetDataSource();
            }

            if (this.p6spy) {
                Field realDataSourceField = P6DataSource.class.getDeclaredField("realDataSource");
                realDataSourceField.setAccessible(true);
                dataSource = (DataSource)realDataSourceField.get(dataSource);
            }

            Class clazz = dataSource.getClass();

            try {
                Method closeMethod = clazz.getDeclaredMethod("close");
                closeMethod.invoke(dataSource);
            } catch (NoSuchMethodException var6) {
                log.warn("dynamic-datasource close the datasource named [{}] failed,", item.getKey());
            }
        }

        log.info("dynamic-datasource all closed success,bye");
    }

    public void afterPropertiesSet() throws Exception {
        Map<String, DataSource> dataSources = this.provider.loadDataSources();
        Iterator var2 = dataSources.entrySet().iterator();

        while(var2.hasNext()) {
            Map.Entry<String, DataSource> dsItem = (Map.Entry)var2.next();
            this.addDataSource((String)dsItem.getKey(), (DataSource)dsItem.getValue());
        }

        if (this.groupDataSources.containsKey(this.primary)) {
            log.info("dynamic-datasource initial loaded [{}] datasource,primary group datasource named [{}]", dataSources.size(), this.primary);
        } else {
            if (!this.dataSourceMap.containsKey(this.primary)) {
                throw new RuntimeException("dynamic-datasource Please check the setting of primary");
            }

            log.info("dynamic-datasource initial loaded [{}] datasource,primary datasource named [{}]", dataSources.size(), this.primary);
        }

    }

    public void setProvider(DynamicDataSourceProvider provider) {
        this.provider = provider;
    }

    public void setStrategy(Class<? extends DynamicDataSourceStrategy> strategy) {
        this.strategy = strategy;
    }

    public void setPrimary(String primary) {
        this.primary = primary;
    }

    public void setStrict(boolean strict) {
        this.strict = strict;
    }

    public boolean getP6spy() {
        return p6spy;
    }

    public boolean getSeata() {
        return seata;
    }
}
